#!/usr/bin/env python
# -*- coding:utf-8 -*-

from unicorn import *
from unicorn.x86_const import *
from capstone import *
import struct
import traceback
from emuapi.windows.winemu import WinEmu
from emuapi.windows.hook import test
from emuapi.asm.asmcode import disasm,get_hex_code,get_local_arch
from emuapi.asm.instruction import optCode
import emuapi.settings as settings
from lib.utils import printer

def pe_emu_51testing():
    path = settings.PE['hello']
    data = open(path, "rb").read()
    tt = WinEmu(data,ispe=True,isteb=False)
    try:
        printer.info("虚拟环境加载的程序地址: " + path)
        tt.build_emu_envir()
        tt.mu.hook_add(UC_HOOK_CODE,test.pe_hook_code,user_data=tt)
        tt.mu.hook_add(UC_HOOK_MEM_INVALID, test.hook_mem_invalid)
        tt.mu.emu_start(0x40141e,0x40141e + 0x1000)
        tt.reg.print_registers()
    except Exception as e:
        printer.warn(str(e))
        tt.reg.print_registers()

def shellcode_emu_51testing(shellcode = 'X86_CODE32'):
    try:
        if shellcode in settings.shellcode.keys():
            binCode = settings.shellcode[shellcode]
        else:
            binCode = shellcode
        dlls = ['ntdll.dll','kernel32.dll','user32.dll']
        tt = WinEmu(binCode,lDll=dlls)
        tt.build_emu_envir()
        tt.mu.hook_add(UC_HOOK_CODE,test.code_hook_code,user_data=tt)
        tt.mu.hook_add(UC_HOOK_MEM_INVALID, test.hook_mem_invalid)
        tt.mu.hook_add(UC_HOOK_INTR,test.code_hook_interrupt)
        tt.mu.emu_start(tt.codeBase,tt.codeBase + len(binCode))
        tt.reg.print_registers()
    except Exception as e:
        print(e)
        traceback.print_exc()
        tt.reg.print_registers()

def generate_binary_code(asmcode:list):
    code = '\n'.join(asmcode)
    hexCode = get_hex_code(code)
    return hexCode

def disam_shellcode_51testing(shellcode,arch):
    print(shellcode)
    if shellcode in settings.shellcode.keys():
        disasm(settings.shellcode[shellcode],arch)
    else:
        print("反汇编内容为空!")

def asmcode_emu(asm:list):
    code = '\n'.join(asm)
    hexCode = get_hex_code(code)
    shellcode_emu_51testing(hexCode)

def get_asm_explain(code):
    explain = ""
    for k,v in optCode.items():
        if code == k:
            explain = k + ": " + v
            break
    return explain
def show_shellcode_choice():
    for k in settings.shellcode.keys():
        print(k)


